#include <iostream>
#include <sstream>
#include <fstream>
#include <cstring>
#include <vector>
#include <queue>
#include <set>
#include <algorithm>
#include <random>
#include "getopt.h"
#include <sys/stat.h>
#include "config-fnt.h"
#include "graph.priv.h"
#include "stream.h"
#include "public/utils.h"

/* Options */
static const char* opts=":o:s:rf:hv";
static const struct option opts_long[]=
{
	{"outdir", 1, nullptr, 'o'},
	{"size", 1, nullptr, 's'},
	{"random", 0, nullptr, 'r'},
	{"file", 1, nullptr, 'f'},
	{"help", 0, nullptr, 'h'},
	{"version", 0, nullptr, 'v'},
	{nullptr, 0, nullptr, 0}
};

/* Print version information and exit */
void version(const char* name) {
	std::cout<<name<<" (" PACKAGE_NAME_LONG ") " PACKAGE_VERSION "\n";
	std::cout<<"Copyright (C) 2015-2016 " PACKAGE_ORG "\n";
	std::cout<<"License GPLv3+: GNU GPL version 3 or later <http://gnu.org/licenses/gpl.html>\n\n";
	std::cout<<"Written by: " PACKAGE_AUTHORS "\n";
	exit(0);
}
// //
/* Print usage information and exit */
void usage(const char* name, int ecode) {
	if(ecode) {
		std::cerr<<"Try '"<<name<<" --help' for more information.\n";
		exit(ecode);
	}
	std::cout<<"Usage: "<<name<<" [OPTION]... FNT_FILE\n";
	std::cout<<
		"Split a FNT file and find errors.\n\n"
		"Options:\n"
		"  -o, --outdir=DIR  [OPTIONAL] Save splitted files to DIR. To current\n"
		"                       working directory by default.\n\n"
		"  -s, --size=N      [OPTIONAL] Splite neurons into groups of size N.\n"
		"                       By default, 1 neuron per group.\n"
		"  -r, --random      [OPTIONAL] Shuffle randomly before spliting.\n"
		"  -f, --file=FILE   [OPTIONAL] Read grouping information from FILE.\n"
		"                       Each line in FILE is the group name of the\n"
		"                       corresponding neuron.\n\n"
		"  -h, --help        [OPTIONAL] Print this help, and do nothing else.\n"
		"  -v, --version     [OPTIONAL] Print version information and exit.\n\n"
		"Report bugs to <" PACKAGE_BUGREPORT ">\n"
		PACKAGE_NAME_LONG " home page: <" PACKAGE_URL ">\n";
	exit(ecode);
}

/* Ensure the dir exist. Try create if not. Return true on success. */
bool ensureDir(const char* dir) {
	struct stat buf;
	int r=stat(dir, &buf);
	if(r!=-1 && S_ISDIR(buf.st_mode)) {
		return true;
	}
#ifdef BUILD_MINGW
	r=mkdir(dir);
#else
	r=mkdir(dir, 0777);
#endif
	if(r!=-1)
		return true;
	std::cerr<<"Cannot create dir '"<<dir<<"': "<<strerror(errno)<<".\n";
	return false;
}

Graph load_graph(const char* filename, std::vector<std::string>& header) {
	Graph gr=GraphPriv::alloc();
	try {
		InputStream fs;
		if(!fs.open(filename))
			throwError("Failed to open file.");
		fs.precision(13);
		fs.setf(std::ios::scientific);

		std::string line;
		fnt_getline(fs, line);
		if(!fs)
			throwError("Failed to read magic header");
		bool gzipped=false;
		if(line==FNTZ_MAGIC) {
			gzipped=true;
		} else if(line!=FNT_MAGIC) {
			throwError("Wrong magic: ", line);
		}
		header.push_back(line);

		fnt_getline(fs, line);
		if(!fs)
			throwError("Failed to read URL line");
		header.push_back(line);

		fnt_getline(fs, line);
		if(!fs)
			throwError("Failed to read line");
		while(line!="BEGIN_TRACING_DATA") {
			header.push_back(line);
			fnt_getline(fs, line);
			if(!fs)
				throwError("Failed to read URL line");
		}

		if(gzipped) {
			if(!fs.pushGzip())
				throwError("Failed to start decompressing");
		}
		auto gp=GraphPriv::get(gr);
		gp->load(fs);
		if(gzipped) {
			if(!fs.pop())
				throwError("Failed to finish decompressing");
		}
		if(!fs.close())
			throwError("Failed to close file.");
		gp->updateModel();
	} catch(const std::exception& e) {
		GraphPriv::free(gr);
	}
	return gr;
}
bool save_graph(Graph gr, const std::string& filename, const std::vector<std::string>& header) {
	bool ret=true;
	try {
		OutputStream fs;
		if(!fs.open(filename))
			throwError("Failed to open file");
		fs.precision(13);
		fs.setf(std::ios::scientific);

		bool gzipped=false;
		if(header[0]==FNTZ_MAGIC)
			gzipped=true;

		const std::vector<std::string> stringsToSkip{
			"core.current.",
				"core.target."
		};
		for(auto& l: header) {
			bool skip=false;
			for(auto& s: stringsToSkip) {
				if(l.compare(0, s.length(), s)==0) {
					skip=true;
					break;
				}
			}
			if(skip)
				continue;
			fs<<l<<'\n';
			if(!fs)
				throwError("Failed to write line");
		}

		fs<<"BEGIN_TRACING_DATA"<<'\n';
		if(!fs)
			throwError("Failed to write data start mark");

		if(gzipped) {
			if(!fs.pushGzip())
				throwError("Failed to start compressing");
		}
		GraphPriv::get(gr)->save(fs);
		if(gzipped) {
			if(!fs.pop())
				throwError("Failed to finish compressing");
		}
		if(!fs.close())
			throwError("Failed to close file");
	} catch(const std::exception& e) {
		ret=false;
	}
	return ret;
}

static int cmp_point(const Point& a, const Point& b) {
	if(a._x<b._x)
		return -1;
	if(a._x>b._x)
		return 1;
	if(a._y<b._y)
		return -1;
	if(a._y>b._y)
		return 1;
	if(a._z<b._z)
		return -1;
	if(a._z>b._z)
		return 1;
	if(a._r<b._r)
		return -1;
	if(a._r>b._r)
		return 1;
	if(a.m<b.m)
		return -1;
	if(a.m>b.m)
		return 1;
	return 0;
}
static int cmp_edge(Edge a, Edge b) {
	auto epa=EdgePriv::get(a);
	auto epb=EdgePriv::get(b);
	if(epa->type<epb->type)
		return -1;
	if(epa->type>epb->type)
		return 1;
	if(epa->points.size()<epb->points.size())
		return -1;
	if(epa->points.size()>epb->points.size())
		return 1;
	auto n=epa->points.size();
	for(size_t i=0; i<n; i++) {
		auto r=cmp_point(epa->points[i], epb->points[i]);
		if(r<0)
			return -1;
		if(r>0)
			return 1;
	}
	return 0;
}
static int cmp_edge_idx(Edge a, Edge b) {
	auto epa=EdgePriv::get(a);
	auto epb=EdgePriv::get(b);
	if(epa->index<epb->index)
		return -1;
	if(epa->index>epb->index)
		return 1;
	return 0;
}
typedef std::pair<Edge, bool> EdgeEnd;
static int cmp_edge_end(const EdgeEnd& a, const EdgeEnd& b) {
	auto r=cmp_edge_idx(a.first, b.first);
	if(r<0)
		return -1;
	if(r>0)
		return 1;
	if(!a.second && b.second)
		return -1;
	if(a.second && !b.second)
		return 1;
	return 0;
}
static int cmp_vert(Vertex a, Vertex b) {
	auto vpa=VertexPriv::get(a);
	auto vpb=VertexPriv::get(b);
	if(!vpa->finished && vpb->finished)
		return -1;
	if(vpa->finished && !vpb->finished)
		return 1;

	if(vpa->neighbors.size()<vpb->neighbors.size())
		return -1;
	if(vpa->neighbors.size()>vpb->neighbors.size())
		return 1;
	auto n=vpa->neighbors.size();
	for(size_t i=0; i<n; i++) {
		auto r=cmp_edge_end(vpa->neighbors[i], vpb->neighbors[i]);
		if(r<0)
			return -1;
		if(r>0)
			return 1;
	}
	return 0;
}
void reorder_graph(GraphPriv* gp) {
	std::sort(gp->edges.begin(), gp->edges.end(), [](Edge a, Edge b)->bool { return cmp_edge(a, b)<0; });
	for(size_t i=0; i<gp->edges.size(); i++) {
		auto ep=EdgePriv::get(gp->edges[i]);
		ep->index=i;
	}
	for(size_t i=0; i<gp->vertices.size(); i++) {
		auto vp=VertexPriv::get(gp->vertices[i]);
		std::sort(vp->neighbors.begin(), vp->neighbors.end(), [](const EdgeEnd& a, const EdgeEnd& b)->bool { return cmp_edge_end(a, b)<0; });
	}
	std::sort(gp->vertices.begin(), gp->vertices.end(), [](Vertex a, Vertex b)->bool { return cmp_vert(a, b)<0; });
}

struct Component {
	std::vector<size_t> edges;
	std::vector<size_t> verts;
	std::vector<size_t> nrns;
};

std::vector<Component> find_components(Graph gr, const std::vector<std::string>& id2group) {
	std::vector<Component> comps;
	auto n=gr.vertices().size();
	std::vector<bool> edg_dirty(gr.edges().size(), false);
	std::vector<bool> vert_dirty(n, false);
	std::vector<bool> nrn_dirty(gr.trees().size(), false);
	while(true) {
		size_t seed=SIZE_MAX;
		for(size_t i=0; i<n; i++) {
			if(!vert_dirty[i]) {
				seed=i;
				break;
			}
		}
		if(seed==SIZE_MAX)
			break;
		Component comp{{}, {}, {}};
		std::queue<size_t> que{};
		que.push(seed);
		while(!que.empty()) {
			auto i=que.front();
			que.pop();
			auto v=gr.vertices()[i];
			if(!vert_dirty[i]) {
				comp.verts.emplace_back(i);
				for(auto& p: v.neighbors()) {
					auto v1=p.second?p.first.leftVertex():p.first.rightVertex();
					que.push(v1.index());
				}
				auto nrn=v.tree();
				if(nrn && nrn.root()==v) {
					for(size_t j=0; j<id2group.size(); j++) {
						if(id2group[j]==id2group[nrn.index()]) {
							que.push(gr.trees()[j].root().index());
						}
					}
				}
				vert_dirty[i]=true;
			}
		}
		for(auto i: comp.verts) {
			auto v=gr.vertices()[i];
			for(auto& p: v.neighbors()) {
				auto j=p.first.index();
				if(!edg_dirty[j]) {
					comp.edges.push_back(j);
					edg_dirty[j]=true;
				}
			}
		}
		for(auto i: comp.verts) {
			auto nrn=gr.vertices()[i].tree();
			if(nrn) {
				auto j=nrn.index();
				if(!nrn_dirty[j]) {
					comp.nrns.push_back(j);
					nrn_dirty[j]=true;
				}
			}
		}
		comps.push_back(std::move(comp));
	}
	for(auto v: edg_dirty) {
		if(!v)
			throwError("Untouched edge");
	}
	for(auto v: vert_dirty) {
		if(!v)
			throwError("Untouched vertex");
	}
	for(auto v: nrn_dirty) {
		if(!v)
			throwError("Untouched neuron");
	}
	return comps;
}


int main(int argc, char* argv[]) {
	const char* outdir=".";

	size_t outsize=1;
	bool shuffle=false;
	const char* groupfile=nullptr;

	// Parse options
	int opt;
	int prev_optind=optind;
	while((opt=getopt_long(argc, argv, opts, opts_long, nullptr))!=-1) {
		switch(opt) {
			case '?':
				std::cerr<<argv[0]<<": unrecognized option '";
				std::cerr<<argv[prev_optind]<<"'\n";
				usage(argv[0], -1);
			case ':':
				std::cerr<<argv[0]<<": option '"<<argv[prev_optind];
				std::cerr<<"' requires an argument\n";
				usage(argv[0], -1);
			case 'h':
				usage(argv[0], 0);
			case 'v':
				version(argv[0]);
			case 'o':
				outdir=optarg;
				break;
			case 's':
				outsize=atol(optarg);
				break;
			case 'r':
				shuffle=true;
				break;
			case 'f':
				groupfile=optarg;
				break;
			default:
				std::cerr<<"Unexpected error.\n";
				usage(argv[0], -1);
		}
		prev_optind=optind;
	}

	if(optind+1<argc) {
		std::cerr<<"Too many input FNT files.\n";
		usage(argv[0], -1);
	}
	if(optind>=argc) {
		std::cerr<<"No input FNT file specified.\n";
		usage(argv[0], -1);
	}

	// Ensure arguments are correct
	if(!ensureDir(outdir)) {
		return -1;
	}

	if(outsize<=0) {
		std::cerr<<"SIZE should be positive.\n";
		return -1;
	}

	std::vector<std::string> header;
	auto graph=load_graph(argv[optind], header);
	if(!graph) {
		std::cerr<<"Failed to load FNT file.\n";
		return -1;
	}
	if(0) {
		auto gp=GraphPriv::get(graph);
		std::default_random_engine rng{};
		std::shuffle(gp->edges.begin(), gp->edges.end(), rng);
		std::shuffle(gp->vertices.begin(), gp->vertices.end(), rng);
		std::shuffle(gp->trees.begin(), gp->trees.end(), rng);
		gp->updateModel();
	}

	std::vector<std::string> id2group{};
	if(groupfile) {
		std::ifstream fgroup{groupfile};
		if(!fgroup)
			throwError("Failed to open group file.");
		std::string line;
		while(fnt_getline(fgroup, line)) {
			if(!line.empty()) {
				id2group.emplace_back(line);
			}
		}
		if(!fgroup.eof())
			throwError("Failed to get line.");
		if(id2group.size()<graph.trees().size()) {
			size_t nother=graph.trees().size()-id2group.size();
			for(size_t i=0; i<nother; i++)
				id2group.emplace_back("OTHER");
		}
	} else {
		std::vector<size_t> mapping(graph.trees().size());
		for(size_t i=0; i<mapping.size(); i++)
			mapping[i]=i;
		if(shuffle)
			std::shuffle(mapping.begin(), mapping.end(), std::default_random_engine());
		for(size_t i=0; i<mapping.size(); i++) {
			std::ostringstream oss{};
			oss.fill('0');
			oss.width(4);
			oss<<mapping[i]/outsize;
			id2group.emplace_back(oss.str());
		}
	}

	size_t junkid=0;
	auto comps=find_components(graph, id2group);
	for(auto& comp: comps) {
		std::vector<std::string> names{};
		for(auto i: comp.nrns) {
			names.push_back(id2group[i]);
		}
		std::sort(names.begin(), names.end());
		auto ite=std::unique(names.begin(), names.end());
		if(ite==names.begin()) {
			Graph grjunk=GraphPriv::alloc();
			auto gpjunk=GraphPriv::get(grjunk);

			for(auto i: comp.edges) {
				gpjunk->edges.push_back(graph.edges()[i]);
			}
			for(auto i: comp.verts) {
				gpjunk->vertices.push_back(graph.vertices()[i]);
			}
			for(auto i: comp.nrns) {
				gpjunk->trees.push_back(graph.trees()[i]);
			}

			if(grjunk.edges().size()>0) {
				reorder_graph(gpjunk);
				gpjunk->updateModel();
				std::ostringstream ossjunk;
				ossjunk<<outdir<<'/';
				ossjunk<<"JUNK_";
				ossjunk.fill('0');
				ossjunk.width(4);
				ossjunk<<junkid<<".fnt";
				auto fnjunk=ossjunk.str();
				bool ret=save_graph(grjunk, fnjunk, header);
				if(!ret) {
					std::cerr<<"Failed to save: "<<fnjunk<<"\n";
				}
			}
			GraphPriv::free(grjunk);
			junkid++;
		} else {
			std::ostringstream oss;
			oss<<outdir<<'/';
			for(auto i=names.begin(); i!=ite; i++) {
				if(i!=names.begin()) {
					oss<<'_';
				}
				oss<<*i;
			}
			oss<<".fnt";
			auto grout=GraphPriv::alloc();
			auto gp=GraphPriv::get(grout);
			for(auto i: comp.edges) {
				gp->edges.push_back(graph.edges()[i]);
			}
			for(auto i: comp.verts) {
				gp->vertices.push_back(graph.vertices()[i]);
			}
			for(auto i: comp.nrns) {
				gp->trees.push_back(graph.trees()[i]);
			}
			reorder_graph(gp);
			gp->updateModel();
			auto fn=oss.str();
			bool ret=save_graph(grout, fn, header);
			if(!ret) {
				std::cerr<<"Failed to save: "<<fn<<"\n";
			}
			GraphPriv::free(grout);
		}
	}


	GraphPriv::free(graph);
	return 0;
}

